import pandas as pd
from typing import List, Dict, Union, Optional
from dbutils.pooled_db import PooledDB
import pymysql
import logging
from retry import retry
from utils.ConfigUtils import read_config


class DorisQueryTool:
    """
    Doris 数据库查询工具类
    功能：
    1. 使用连接池管理数据库连接
    2. 自动重试失败的查询
    3. 返回 Pandas DataFrame 或原始字典列表
    4. 支持上下文管理（with语句）
    """

    def __init__(
            self,
            pool_size: int = 5,
    ):
        dorisConfig = read_config("E:\\workspace\\pythonWorkspace\\flow-prediction\\config\\jdbc.properties")
        #dorisConfig = read_config("/home/running/app/python/flow-prediction/config/jdbc.properties")
        self.host = dorisConfig.get("doris_host")
        self.port = int(dorisConfig.get("doris_port"))
        self.user = dorisConfig.get("doris_user")
        self.password = dorisConfig.get("doris_password")
        self.database = dorisConfig.get("doris_db")
        self.pool = self._create_pool(pool_size)

    def _create_pool(self, pool_size: int) -> PooledDB:
        """创建Doris连接池"""
        return PooledDB(
            creator=pymysql,
            host=self.host,
            port=self.port,
            user=self.user,
            password=self.password,
            database=self.database,
            maxconnections=pool_size,
            ping=1,  # 自动检测连接有效性
        )

    @retry(tries=3, delay=1, backoff=2, logger=None)
    def execute_query(
            self,
            sql: str,
            params: Optional[Union[Dict, List]] = None,
            return_df: bool = True,
    ) -> Union[pd.DataFrame, List[Dict]]:
        """
        执行SQL查询
        :param sql: 要执行的SQL语句
        :param params: 查询参数（可选）
        :param return_df: 是否返回DataFrame（默认True）
        :return: DataFrame 或 字典列表
        """
        conn = self.pool.connection()
        try:
            with conn.cursor(pymysql.cursors.DictCursor) as cursor:
                cursor.execute(sql, params)
                result = cursor.fetchall()

                if return_df and result:
                    return pd.DataFrame(result)
                return result
        except Exception as e:
            logging.error(f"执行SQL失败: {e}\nSQL: {sql}")
            raise
        finally:
            conn.close()

    def __enter__(self):
        """支持with上下文管理"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """退出时关闭连接池"""
        self.pool.close()


# 使用示例
if __name__ == "__main__":
    # 初始化配置（建议从配置文件读取）
    # 使用示例
    with DorisQueryTool() as tool:
        # 示例1：返回DataFrame
        df = tool.execute_query("select count(1) FROM toll_ex where EX_TIME >= '2025-1-1 00:00:00';")
        print(f"查询结果（DataFrame格式）:\n{df.head()}")

